{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57535272",
   "metadata": {},
   "source": [
    "### 导入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05b89b82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fde7370",
   "metadata": {},
   "source": [
    "### 读取文本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970d2900",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../data/10027.txt\", \"r\", encoding=\"utf-8\") as file:\n",
    "    text = file.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "297748a5",
   "metadata": {},
   "source": [
    "### 分块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "200f9793",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1️⃣ 分块器初始化\n",
    "splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=256,\n",
    "    chunk_overlap=50,\n",
    "    separators=[\"\\n\\n\", \"\\n\", \"。\", \"；\", \"？\", \"！\", \".\", \"!\", \"?\", \"，\", \" \", \"\"]\n",
    ")\n",
    "chunks = splitter.split_text(text)\n",
    "print(type(chunks))\n",
    "\n",
    "# 保存分块结果到 txt 文件\n",
    "with open(\"../data/chunks_output.txt\", \"w\", encoding=\"utf-8\") as f:\n",
    "    for i, chunk in enumerate(chunks):\n",
    "        f.write(f\"[Chunk {i+1}]: {chunk}\\n\\n\")\n",
    "\n",
    "# print(f\"📌 共分出 {len(clean_chunks)} 个 chunk：\")\n",
    "# for i, chunk in enumerate(clean_chunks):\n",
    "#     print(f\"[Chunk {i+1}]: {chunk}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f49fda6c",
   "metadata": {},
   "source": [
    "### 编码（bge-small-en-v1.5）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f378452",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SentenceTransformer(\"../model/bge-small-en-v1.5\", device=\"cuda\")\n",
    "vectors = model.encode(chunks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46552845",
   "metadata": {},
   "source": [
    "### 在docker中运行Weaviate\n",
    "\n",
    "#### 下载Weaviate镜像\n",
    "\n",
    "~~~\n",
    "docker pull semitechnologies/weaviate\n",
    "~~~\n",
    "\n",
    "#### 在docker中启动weaviate\n",
    "\n",
    "~~~\n",
    "docker run -p 8080:8080\\\n",
    "    -v $(pwd)/weaviate_storage:/weaviate/storage \\\n",
    "    --env DISABLE_GRPC=true \\\n",
    "    semitechnologies/weaviate\n",
    "~~~"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "340f0ed8",
   "metadata": {},
   "source": [
    "### 为Weaviate创建一个客户端对象"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba57e05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from weaviate.client import WeaviateClient\n",
    "from weaviate.connect import ConnectionParams\n",
    "\n",
    "\n",
    "# 初始化Weaviate客户端\n",
    "client = WeaviateClient(\n",
    "    connection_params=ConnectionParams.from_url(\n",
    "        url=\"http://172.20.50.49:8080\",  # REST 服务地址\n",
    "        grpc_port=50051  # 可以保留这个参数，虽然不用 gRPC 也无所谓\n",
    "    ),\n",
    "    skip_init_checks=True  # ✅ 关键参数：跳过连接时的 gRPC 检查\n",
    ")\n",
    "\n",
    "client.connect()\n",
    "\n",
    "# 检查连接是否成功\n",
    "if client.is_ready():\n",
    "    print(\"Weaviate is ready\")\n",
    "else:\n",
    "    print(\"Weaviate is not ready\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ce10926",
   "metadata": {},
   "source": [
    "### 检测集合是否存在"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f30a4852",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_collection_exists(client: weaviate.WeaviateClient, collection_name: str) -> bool:\n",
    "    \"\"\"\n",
    "    检查集合是否存在\n",
    "    :param client: Weaviate 客户端\n",
    "    :param collection_name: 集合名称\n",
    "    :return: True 或 False\n",
    "    \"\"\"\n",
    "    try:\n",
    "        collections = client.collections.list_all()\n",
    "        return collection_name in collections\n",
    "    except Exception as e:\n",
    "        print(f\"检查集合异常: {e}\")\n",
    "        return False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1923ab06",
   "metadata": {},
   "source": [
    "### 创建集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b625e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_collection(client: weaviate.WeaviateClient, collection_name: str):\n",
    "    \"\"\"\n",
    "    创建集合\n",
    "    :param client: Weaviate 客户端\n",
    "    :param collection_name: 集合名称\n",
    "    \"\"\"\n",
    "    collection_obj = {\n",
    "        \"class\": collection_name,\n",
    "        \"description\": \"A collection for product information\",  # 描述集合的用途\n",
    "        \"vectorizer\": \"none\",   # 使用外部向量化\n",
    "        \"vectorIndexType\": \"hnsw\",  # 向量索引类型\n",
    "        \"vectorIndexConfig\": {\n",
    "            \"distance\": \"cosine\",   # 向量距离度量方式\n",
    "            \"efConstruction\": 200,  # HNSW 构建时的 ef 值\n",
    "            \"maxConnections\": 64    # HNSW 最大连接数\n",
    "        },\n",
    "        \"properties\": [\n",
    "            {\n",
    "                \"name\": \"text\", # 属性名称\n",
    "                \"description\": \"The text content\",  # 属性描述\n",
    "                \"dataType\": [\"text\"],   # 属性数据类型\n",
    "                \"tokenization\": \"None\", # 分词方式\n",
    "                \"indexFilterable\": True,    # 是否可过滤\n",
    "                \"indexSearchable\": True # 是否可搜索\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "    try:\n",
    "        client.collections.create_from_dict(collection_obj)\n",
    "        print(f\"创建集合 '{collection_name}' 成功.\")\n",
    "    except weaviate.exceptions.UnexpectedStatusCodeException as e:\n",
    "        print(f\"创建集合异常: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeff9f16",
   "metadata": {},
   "source": [
    "### 插入数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f0f326",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_datas(client: weaviate.WeaviateClient, collection_name: str, datas: list):\n",
    "    \"\"\"\n",
    "    向集合中插入数据\n",
    "    :param client: Weaviate 客户端\n",
    "    :param collection_name: 集合名称\n",
    "    :param datas: 数据列表\n",
    "    \"\"\"\n",
    "    collection = client.collections.get(collection_name)\n",
    "    for i in range(len(datas)):\n",
    "        item = datas[i]\n",
    "        text = item[\"text\"]\n",
    "        vector = item[\"vector\"]\n",
    "        properties = {\"text\": text}\n",
    "        try:\n",
    "            uuid = collection.data.insert(properties=properties, vector=vector)\n",
    "            print(f\"文档添加内容: {text[:30]}..., uuid: {uuid}\")\n",
    "        except Exception as e:\n",
    "            print(f\"添加文档异常: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e30b93d",
   "metadata": {},
   "source": [
    "### 查询数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0951383c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_vector_collection(client: weaviate.WeaviateClient, collection_name: str, query: str, k: int) -> list:\n",
    "    \"\"\"\n",
    "    从集合中查询数据\n",
    "    :param client: Weaviate 客户端\n",
    "    :param collection_name: 集合名称\n",
    "    :param query: 查询字符串\n",
    "    :param k: 返回的结果数量\n",
    "    :return: 查询结果列表\n",
    "    \"\"\"\n",
    "    vector = model.encode(query)  # 将查询字符串转换为向量\n",
    "    collection = client.collections.get(collection_name)\n",
    "    response = collection.query.near_vector(\n",
    "        near_vector=vector,\n",
    "        limit=k\n",
    "    )\n",
    "    documents = [res.properties['text'] for res in response.objects]\n",
    "    return documents\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b7dc92d",
   "metadata": {},
   "source": [
    "### 构建数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c42a93b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectors = [[ np.random.uniform(-1, 1) for _ in range(384) ] for _ in range(len(chunks)) ]\n",
    "datas = [ {\"id\": i, \"vector\": vectors[i], \"text\": chunks[i]} for i in range(len(vectors)) ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53acd3ff",
   "metadata": {},
   "source": [
    "### 检查集合是否存在"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd4f990d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not check_collection_exists(client, \"E10027\"):\n",
    "    # 创建集合\n",
    "    create_collection(client, \"E10027\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37c22240",
   "metadata": {},
   "source": [
    "### 插入数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c2d7a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_datas(client, \"E10027\", datas)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e89f91fd",
   "metadata": {},
   "source": [
    "### 查询数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b6b39d",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results = query_vector_collection(client, \"E10027\", \"工作时间\", 5)\n",
    "print(\"查询结果:\", query_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3b59258",
   "metadata": {},
   "source": [
    "### 查询数据有时候会报错，报错原因：\n",
    "- Weaviate 服务未启动：检查一下目标机器上的 Weaviate 服务是否正在运行。你可以尝试在该服务器上执行 docker ps（如果使用 Docker 部署的话）来检查 Weaviate 容器的状态，或者使用其他方式确认 Weaviate 服务是否正在监听端口 50051。（当时正在运行，远程服务器也在监听50051端口）\n",
    "\n",
    "- 网络问题：如果目标服务器的网络有问题（如防火墙设置、网络断开等），就会导致无法连接。确认你的机器能够访问 172.20.50.49 这个 IP 地址。你可以使用 ping 172.20.50.49 来检查网络连接，或尝试用 telnet 172.20.50.49 50051 测试是否能成功连接到该端口。（能ping通）\n",
    "\n",
    "- 端口被阻塞：确认端口 50051 没有被防火墙或者安全组规则阻塞。如果你使用云服务（如 AWS、Azure 等），请检查相关的安全组或防火墙配置，确保 50051 端口是开放的。\n",
    "\n",
    "- Weaviate 配置问题：检查 Weaviate 的配置文件，确保服务配置正确并允许外部连接。特别是检查是否正确绑定了服务器的 IP 地址和端口。\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
